/*
 * Copyright 2006-2007 S W McIlroy
 * Licensed  under the  Apache License,  Version 2.0  (the "License");
 * you may not use  this file  except in  compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed  under the  License is distributed on an "AS IS" BASIS,
 * WITHOUT  WARRANTIES OR CONDITIONS  OF ANY KIND, either  express  or
 * implied.
 *
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.mockfu;

import static org.objectweb.asm.Opcodes.*;

import java.lang.reflect.Method;

import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Type;

/**
 * A method replacement specifies a method invocation for a {@link UnitBuilder} to replace on a
 * class it generates. The invocation may be on an object, on a constructor, or static.
 *
 * @param <U> the type of the unit class
 * @param <C> the type of the method's return value
 */
public class MethodReplacement<U, C> extends Replacement<U>
{
    private Object replacementValue;
    private final ConstrainedMethod method;

    MethodReplacement(Class<U> unitClass, Class<C> ctorClass, Object... args)
    {
        super(unitClass);
        this.method = new ConstrainedMethod(ctorClass, args);
    }

    MethodReplacement(Class<U> unitClass, ConstrainedMethod method)
    {
        super(unitClass);
        this.method = method;
    }

    /**
     * Sets a stub value to replace the method invocation with.
     *
     * @param replacement
     * @return
     */
    public MethodReplacement<U, C> with(C replacement)
    {
        if (replacement instanceof MockControl)
            this.replacementValue = ((MockControl<?>)replacement).mock();
        else
            this.replacementValue = replacement;
        return this;
    }


    /**
     * Specifies a method on the unit within which replacements may occur. If no methods are
     * specified replacements occur on all methods.
     */
    @Override
    public MethodReplacement<U, C> in(Method unitMethod)
    {
        super.in(unitMethod);
        return this;
    }

    @Override
    Object getValue()
    {
        return replacementValue;
    }

    @Override
    boolean matches(Object context)
    {
        return method.matches((Object[])context);
    }

    @Override
    ClassVisitor chain(ClassVisitor visitor)
    {
        return new ClassAdapter(visitor) {
            @Override
            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions)
            {
                MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
                if (mv != null)
                {
                    if (methods.isEmpty())
                    {
                        return new ReplacementMethodVisitor(mv, access, name, desc);
                    }
                    else
                    {
                        for (Method m : methods)
                            if (m.getName().equals(name) && desc.equals(Type.getMethodDescriptor(m)))
                                return new ReplacementMethodVisitor(mv, access, name, desc);
                    }
                }
                return mv;
            }

        };
    }

    private class ReplacementMethodVisitor extends NormalMethodGenerator
    {

        public ReplacementMethodVisitor(MethodVisitor delegate, int access, String name, String desc)
        {
            super(delegate, access, name, desc);
        }

        @Override
        public void visitTypeInsn(int opcode, String desc)
        {
            if (opcode == NEW && desc.equals(Type.getInternalName(method.getDeclaringClass())))
            {
                mv = new ConstructorReplacementVisitor(mv);
            }
            mv.visitTypeInsn(opcode, desc);
        }

        @Override
        public void visitMethodInsn(int opcode, String owner, String name, String desc)
        {
            Type methodType = Type.getType(method.getDeclaringClass());

            if (name.equals(method.getName()) && owner.equals(methodType.getInternalName()))
            {
                Type[] argumentTypes = Type.getArgumentTypes(desc);
                if (method.isMethod(name, argumentTypes))
                {
                    ConstructorReplacementVisitor recorder = null;
                    if (name.equals("<init>") && mv instanceof ConstructorReplacementVisitor)
                    {
                        /* stop recording */
                        recorder = (ConstructorReplacementVisitor)mv;
                        mv = recorder.playbackArguments();
                    }
                    Type objectType = Type.getType(Object.class);
                    Type returnType = Type.getReturnType(desc);
                    if (name.equals("<init>"))
                        returnType = methodType;

                    /* create an array with method arguments from top of stack */
                    int arraySlot = storeArray(argumentTypes);

                    /* match replacement object */
                    push(id);
                    loadLocal(arraySlot);
                    invokeStatic("org/mockfu/Replacement", "match", "(ILjava/lang/Object;)Ljava/lang/Object;");

                    Type replacementType = Type.getType(Replacement.class);
                    Label matchedLabel = new Label();
                    Label endMatchedLabel = new Label();

                    /* if replacement value matches magic "NO_MATCH" constant */
                    dup();
                    getStatic(replacementType, "NOT_MATCHED", objectType);
                    ifCmp(objectType, NE, matchedLabel);

                    /* then restore stack and invoke constructor */
                    pop();
                    if (recorder != null)
                        recorder.playbackInstantiation();

                    for (int i = 0 ; i < argumentTypes.length ; i++)
                    {
                        loadLocal(arraySlot);
                        push(i);
                        arrayLoad(objectType);
                        Type type = argumentTypes[i];
                        /* unbox primitive types, or cast object types */
                        unbox(type);
                    }
                    mv.visitMethodInsn(opcode, owner, name, desc);
                    goTo(endMatchedLabel);

                    mark(matchedLabel);

                    if (opcode != INVOKESTATIC && recorder == null)
                    {
                        swap();
                        pop();
                    }
                    /* cast or unbox replacement */
                    unbox(returnType);

                    mark(endMatchedLabel);

                    if (recorder != null)
                        recorder.adjustCopies();

                    return;
                }
                else if (name.equals("<init>") && mv instanceof ConstructorReplacementVisitor)
                {
                    /* we set up a recorder but the constructor didn't match our arguments: play everything back */
                    ConstructorReplacementVisitor recorder = (ConstructorReplacementVisitor)mv;
                    mv = recorder.playback();
                }
            }
            super.visitMethodInsn(opcode, owner, name, desc);
        }

    }

}
